import glob
import numpy as np
import torchio as tio
import logging
import os
import skimage.exposure as exposure
import nibabel as nib
import matplotlib.pyplot as plt


def resample_by_torchio(input_path, output_path):
    # 加载图像
    image = tio.ScalarImage(input_path)

    # 定义变换
    transform = tio.Compose(
        [
            tio.ToCanonical(),  # 将图像转换为标准方向
            tio.RescaleIntensity((0, 255)),  # 重新缩放强度到 [0, 255]
        ]
    )

    # 应用变换
    transformed_image = transform(image)

    # 获取图像数据并转换为 uint8
    img_data = transformed_image.data.numpy()[0]  # 形状 (x, y, z)
    img_data_uint8 = img_data.astype(np.uint8)  # 强制转换为 uint8

    # 使用 nibabel 保存为 NIfTI 格式
    nifti_img = nib.Nifti1Image(img_data_uint8, affine=transformed_image.affine)
    nifti_img.header.set_data_dtype(np.uint8)  # 显式指定数据类型为 uint8
    nib.save(nifti_img, output_path)

    # 调试：检查保存后的范围
    saved_image = tio.ScalarImage(output_path)
    print(f"Saved range: min={saved_image.data.min()}, max={saved_image.data.max()}")


def save_middle_slices(volume, out_dir, name="test_1"):
    """
    Save middle slices along x, y, z axes of a 3D volume.

    :param volume: A 3D numpy array representing the volume (shape should be [depth, height, width]).
    :param output_dir: Directory to save the slice images.
    """
    print(f"Volume shape: {volume.shape}")
    print(f"Volume dtype: {volume.dtype}")
    print(f"Volume range: min={volume.min()}, max={volume.max()}")
    volume = volume.squeeze()
    print(volume.shape)
    # Calculate the middle indices
    mid_x, mid_y, mid_z = (
        volume.shape[1] // 2,
        volume.shape[2] // 2,
        volume.shape[0] // 2,
    )
    # Extract the middle slices along each axis
    slice_x = volume[mid_z, :, :]
    slice_y = volume[:, mid_y, :]
    slice_z = volume[:, :, mid_x]
    # 检查切片的强度分布
    print(
        f"Slice_x range: min={slice_x.min()}, max={slice_x.max()}, mean={slice_x.mean()}"
    )
    print(
        f"Slice_y range: min={slice_y.min()}, max={slice_y.max()}, mean={slice_y.mean()}"
    )
    print(
        f"Slice_z range: min={slice_z.min()}, max={slice_z.max()}, mean={slice_z.mean()}"
    )
    # Save the slices
    plt.figure(figsize=(15, 5))

    # Save Z-axis slice
    plt.subplot(1, 3, 1)
    plt.imshow(slice_x, cmap="gray", vmin=slice_x.min(), vmax=slice_x.max())
    plt.title("Middle Slice Along Z-Axis")
    plt.axis("off")
    # plt.savefig(f"{name}_middle_slice_z.png", bbox_inches='tight')

    # Save Y-axis slice
    plt.subplot(1, 3, 2)
    plt.imshow(slice_y, cmap="gray", vmin=slice_y.min(), vmax=slice_y.max())
    plt.title("Middle Slice Along Y-Axis")
    plt.axis("off")
    # plt.savefig(f"{name}_middle_slice_y.png", bbox_inches='tight')

    # Save X-axis slice
    plt.subplot(1, 3, 3)
    plt.imshow(slice_z, cmap="gray", vmin=slice_z.min(), vmax=slice_z.max())
    plt.title("Middle Slice Along X-Axis")
    plt.axis("off")
    plt.savefig(f"{out_dir}/{name}_middle_slice_.png", bbox_inches="tight")

    plt.close()  # Close the figure to free up memory


def overlay_mask_on_image(
    image_path, mask_path, output_dir="output", slice_axis=0, slice_idx=None
):
    """
    将掩码以红色通道叠加在原始图像上，并显示或保存中间切片。

    Args:
        image_path (str): 原始图像的 .nii.gz 文件路径
        mask_path (str): 掩码的 .nii.gz 文件路径
        output_dir (str): 保存输出图像的目录
        slice_axis (int): 切片轴（0: z轴, 1: y轴, 2: x轴）
        slice_idx (int, optional): 指定切片索引，若为 None 则取中间切片
    """
    # 加载原始图像和掩码
    name = os.path.basename(mask_path).split(".")[0].split("_")[-1]
    img_nii = nib.load(image_path)
    mask_nii = nib.load(mask_path)
    img_data = img_nii.get_fdata()
    mask_data = mask_nii.get_fdata()

    # 确保图像和掩码尺寸一致
    assert img_data.shape == mask_data.shape, "Image and mask dimensions must match!"

    # 归一化原始图像到 [0, 1] 用于显示
    img_data = (img_data - img_data.min()) / (img_data.max() - img_data.min() + 1e-6)

    # 将掩码转换为二值形式（0 或 1）
    mask_data = (mask_data > 0).astype(np.uint8)

    # 选择中间切片（若未指定 slice_idx）
    if slice_idx is None:
        slice_idx = img_data.shape[slice_axis] // 2
    else:
        slice_idx = img_data.shape[slice_axis] // 2 + slice_idx  # 一中间为轴做偏移

    # 根据轴选择切片
    if slice_axis == 0:  # x轴
        img_slice = img_data[slice_idx, :, :]
        mask_slice = mask_data[slice_idx, :, :]
    elif slice_axis == 1:  # y轴
        img_slice = img_data[:, slice_idx, :]
        mask_slice = mask_data[:, slice_idx, :]
    elif slice_axis == 2:  # z轴
        img_slice = img_data[:, :, slice_idx]
        mask_slice = mask_data[:, :, slice_idx]
    else:
        raise ValueError("slice_axis must be 0, 1, or 2")

    # 创建 RGB 图像
    rgb_image = np.zeros((*img_slice.shape, 3), dtype=np.float32)

    # 原始图像作为灰度，填充到 R、G、B 通道
    rgb_image[..., 0] = img_slice  # R
    rgb_image[..., 1] = img_slice  # G
    rgb_image[..., 2] = img_slice  # B

    # 掩码覆盖红色通道
    rgb_image[..., 0] = np.where(mask_slice, 1.0, rgb_image[..., 0])  # 掩码区域设为红色
    rgb_image[..., 1] = np.where(mask_slice, 0.0, rgb_image[..., 1])  # 去除绿色
    rgb_image[..., 2] = np.where(mask_slice, 0.0, rgb_image[..., 2])  # 去除蓝色

    # 显示图像
    plt.figure(figsize=(8, 8))
    plt.imshow(rgb_image)
    plt.title(f"Slice along axis {slice_axis}, index {slice_idx}")
    plt.axis("off")

    # 创建输出目录并保存

    os.makedirs(output_dir, exist_ok=True)
    plt.savefig(
        f"{output_dir}/{name}_axis_{slice_axis}_idx{slice_idx}.png",
        bbox_inches="tight",
    )
    plt.close()

    print(
        f"Overlay image saved to {output_dir}/overlay_slice_axis{slice_axis}_idx{slice_idx}.png"
    )


def resample_unite(dir_list=["./proposed"], target_dir="./processed"):
    """
    将多个目录中的图像统一重采样到一个目标目录。
    """
    subject_id = "infer"  # 统一的target_dir

    os.makedirs(target_dir, exist_ok=True)
    for dir in dir_list:
        input_files = glob.glob(os.path.join(dir, "*.nii.gz"))
        input_files = [
            file
            for file in input_files
            if os.path.basename(file).startswith(f"{subject_id}")
        ]
        logging.info(f"Processing files{input_files}")
        for img_path in input_files:
            output_path = os.path.basename(img_path).replace(
                ".nii.gz", "_registered.nii.gz"
            )
            output_path = os.path.join(target_dir, output_path)
            resample_by_torchio(img_path, output_path)


def resample_unite2(dir_list=["./proposed"], target_dir="./processed"):
    """
    将多个目录中的图像统一重采样到按照规则自动创建的目标目录。
    Params:
        dir_list: 需要处理的目录列表
        target_dir: 统一的目标目录
    """
    for dir in dir_list:
        if os.path.exists(target_dir):
            pass
        else:
            os.makedirs(target_dir, exist_ok=True)
            logging.info(f"Created target directory: {target_dir}")
        input_files = glob.glob(os.path.join(dir, "**", "*.nii.gz"), recursive=True)
        # input_files = [
        #     file
        #     for file in input_files
        #     if os.path.basename(file).startswith(f"{subject_id}")
        # ]
        logging.info(f"Processing files{input_files}")
        for img_path in input_files:
            output_path = os.path.basename(img_path).replace(
                ".nii.gz", "_registered.nii.gz"
            )
            output_path = os.path.join(target_dir, output_path)
            resample_by_torchio(img_path, output_path)


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
    )
    logger = logging.getLogger(__name__)
    # dir_list = ["./B-spline", "./smore", "./mcsr", "./proposed"]

    # if img_path.endswith("_cubic.nii.gz"):
    #     resample_by_torchio(img_path, output_path)
    # else:
    #     resample_by_torchio_r(
    #         img_path,
    #         output_path,
    #         "./B-spline/0_cubic.nii.gz",
    #     )
    # logger.info(f"Resampled image saved to {output_path}")
    # 1. 这里是统一数据范围的
    # dir_list = [
    #     "../prediction/CHEN-REN-GENG/contrast_methods",
    #     "../prediction/CHEN-REN-GENG/source",
    #     "../prediction/CHEN-REN-GENG/proposed/cubic__fix_ds_6x4",
    # ]
    dir_list = [
        "../prediction/DJY/contrast_methods",
        "../prediction/DJY/source",
        "../prediction/DJY/proposed/cubic__fix_ds_6x4",
    ]
    target_dir = "../prediction/DJY/processed"
    resample_unite2(
        dir_list,
        target_dir=target_dir,
    )

    # # 3.这里是将分割结果打在原图上可视化
    # registered_files = [
    #     "./processed/0_ResShift_registered.nii.gz",
    # ]
    # for file in registered_files:
    #     maskname = (
    #         os.path.basename(file)
    #         .replace("0", "prostate")
    #         .replace("_registered.nii.gz", ".nii.gz")
    #     )

    #     mask_path = os.path.join(
    #         "/media/userdisk5/gqyang/final_result_t2map/processed/segmentations", maskname
    #     )
    #     overlay_mask_on_image(
    #         file, mask_path, output_dir="out/sag", slice_axis=0, slice_idx=None
    #     )  # sagittal view 的
    #     for i in range(-30, 30, 1):
    #         overlay_mask_on_image(
    #             file, mask_path, output_dir=f"out/tra{i}", slice_axis=2, slice_idx=i
    #         )

    # # 从这里开始是展示分割效果的
    # registered_files = glob.glob(
    #     os.path.join("/media/userdisk5/gqyang/final_result/processed", "*.nii.gz")
    # )
    # out = "picture"
    # os.makedirs(out, exist_ok=True)
    # for file in registered_files:
    #     name = os.path.basename(file).split(".")[0]
    #     print(name)
    #     image1 = tio.ScalarImage(file)
    #     save_middle_slices(image1.data.numpy(), out_dir=out, name=name)
    # image_path = (
    #     "/media/userdisk5/gqyang/final_result/processed/0_cubic_registered.nii.gz"
    # )

    # mask_path = "/media/userdisk5/gqyang/final_result/processed/segmentations/prostate_cubic.nii.gz"
